import pandas as pd
import torch
from sklearn.metrics import accuracy_score, f1_score
import os

from src.bert.my_dataloader import my_dataLoader
from src.bert.my_model import BertClassifier
from src.bert.stu.stu_bert import StudentBertClassifier
from utils.config import Config

config = Config('E:/Python+AI/group4_nlp_project')


def predict(save_predictions=False):
    # 检查是否存在蒸馏模型，如果存在则使用蒸馏模型，否则使用原始模型
    pruned_model_path = os.path.join(config.save_model, 'bt_model_distilled.pth')
    # pruned_model_path = os.path.join(config.save_model, 'bt_model_pruned.pth')
    # pruned_model_path = os.path.join(config.save_model, 'bt_model_quantized.pth')
    original_model_path = os.path.join(config.save_model, 'bt_model.pth')

    if os.path.exists(pruned_model_path):
        print("使用蒸馏模型进行预测")
        model = StudentBertClassifier(config)
        model.load_state_dict(torch.load(pruned_model_path))
    elif os.path.exists(pruned_model_path):
        print("使用量化模型进行预测")
        # 为量化模型创建一个量化版本的学生模型
        model = StudentBertClassifier(config)
        # 先加载模型到CPU，然后进行量化
        model.load_state_dict(torch.load(pruned_model_path, map_location=torch.device('cpu')))
    elif os.path.exists(pruned_model_path):
        print("使用剪枝模型进行预测")
        model = StudentBertClassifier(config)
        model.load_state_dict(torch.load(pruned_model_path, map_location=torch.device('cpu')))
    elif os.path.exists(original_model_path):
        print("使用原始模型进行预测")
        model = BertClassifier()
        model.load_state_dict(torch.load(original_model_path, map_location=torch.device('cpu')))
    else:
        raise FileNotFoundError("未找到模型文件")

    model.to(config.device)
    model.eval()
    test_loader = my_dataLoader(config.test_path)

    # 加载原始数据以获取文本
    df = pd.read_csv(config.test_path)
    all_texts = df['review_clean'].tolist()

    # 用于存储所有预测和真实标签
    all_cat_preds = []
    all_cat_labels = []
    all_label_preds = []
    all_label_labels = []

    with torch.no_grad():
        for i, (input_ids, attention_mask, labels, cat_labels) in enumerate(test_loader):
            input_ids = input_ids.to(config.device)
            attention_mask = attention_mask.to(config.device)
            labels = labels.float().to(config.device)
            cat_labels = cat_labels.long().to(config.device)

            # 根据模型类型处理不同的输出
            if isinstance(model, StudentBertClassifier):
                out_cat, out_label, _ = model(input_ids, attention_mask)
            else:
                out_cat, out_label = model(input_ids, attention_mask)

            # 分类任务预测
            _, predicted_cat = torch.max(out_cat.data, 1)

            # 标签任务预测（二分类）
            predicted_labels = (torch.sigmoid(out_label).squeeze() > 0.5).float()

            # 收集预测结果和真实标签
            all_cat_preds.extend(predicted_cat.cpu().numpy())
            all_cat_labels.extend(cat_labels.cpu().numpy())
            all_label_preds.extend(predicted_labels.cpu().numpy())
            all_label_labels.extend(labels.cpu().numpy())

    # 使用sklearn计算准确率
    cat_accuracy = accuracy_score(all_cat_labels, all_cat_preds)
    label_accuracy = accuracy_score(all_label_labels, all_label_preds)
    cat_f1 = f1_score(all_cat_labels, all_cat_preds, average='weighted')
    label_f1 = f1_score(all_label_labels, all_label_preds, average='weighted')

    print(f'分类任务F1-score：{cat_f1:.4f}')
    print(f'标签任务F1-score：{label_f1:.4f}')
    print(f'分类任务准确率：{cat_accuracy:.4f}')
    print(f'标签任务准确率：{label_accuracy:.4f}')

    # 如果需要保存预测结果，则保存到CSV文件
    if save_predictions:
        predictions_df = pd.DataFrame({
            'text': all_texts,
            'true_category': all_cat_labels,
            'predicted_category': all_cat_preds,
            'true_label': all_label_labels,
            'predicted_label': all_label_preds
        })
        predictions_file = os.path.join(config.data_path, 'predictions_with_text.csv')
        predictions_df.to_csv(predictions_file, index=False, encoding='utf-8')
        print("预测结果已保存到:", predictions_file)

    return all_texts, all_cat_labels, all_cat_preds, all_label_labels, all_label_preds


if __name__ == '__main__':
    predict()
